#!/usr/bin/env python
#coding:utf-8
import sys
import traceback  

from thrift.transport import TTransport
from cStringIO import StringIO
from struct import pack, unpack
from thrift.Thrift import TException
from lib.ContextSeralize import ContextSeralize
from lib.ContextReader import ContextReader
from lib.JMTBinaryProtocol import JMTBinaryProtocol

class JMTFramedTransport(TTransport.TFramedTransport):

    """Class that wraps another transport and frames its I/O when writing."""

    def __init__(self, trans,):
        self.__trans = trans
        self.__rbuf = StringIO()
        self.__wbuf = StringIO()

    def setContext(self, context):
        self.__context = context

    def isOpen(self):
        return self.__trans.isOpen()

    def open(self):
        return self.__trans.open()

    def close(self):
        return self.__trans.close()

    def readAll(self, sz):
        buff = ''
        have = 0
        while (have < sz):
          chunk = self.read(sz - have)
          have += len(chunk)
          buff += chunk

          if len(chunk) == 0:
            raise EOFError()

        return buff

    def read(self, sz):
        ret = self.__rbuf.read(sz)
        if len(ret) != 0:
          return ret

        bufStr = self.__rbuf.getvalue()
        if (len(bufStr) <= 0) :
            self.readFrame()
            # self.readContext()
        
        return self.__rbuf.read(sz)

    def readContext(self):
        buf = self.__rbuf.read(4)
        sz = unpack('!I', buf)
        if ((sz[0] & JMTBinaryProtocol.VERSION_MASK) != JMTBinaryProtocol.VERSION_1):
            return 0
        # read seqid
        buff = self.__rbuf.read(4)
        seqid, = unpack('!i', buff)
        # read byte1
        buff = self.__rbuf.read(1)
        byte1, = unpack('!b', buff)

        # read byte2
        buff = self.__rbuf.read(1)
        byte2, = unpack('!b', buff)

    def readFrame(self):
        buff = self.__trans.readAll(4)
        sz, = unpack('!i', buff)
        self.__rbuf = StringIO(self.__trans.readAll(sz))

    def write(self, buf):
        self.__wbuf.write(buf)

    def flush(self):
        contextBuf = self.writeContextSerialize()
        wout = self.__wbuf.getvalue()
        # wout += contextBuf
        contextBuf += wout
        wsz = len(contextBuf)
        # reset wbuf before write/flush to preserve state on underlying failure
        self.__wbuf = StringIO()
        # N.B.: Doing this string concatenation is WAY cheaper than making
        # two separate calls to the underlying socket object. Socket writes in
        # Python turn out to be REALLY expensive, but it seems to do a pretty
        # good job of managing string buffer operations without excessive copies
        buf = pack("!i", wsz) + contextBuf
        self.__trans.write(buf)
        self.__trans.flush()

    def writeContextSerialize(self):
        context = self.__context.get()
        if (len(context) <= 0):
            return ''
        obj = ContextSeralize(context)
        transObj = TTransport.TMemoryBuffer()
        protocol = JMTBinaryProtocol(transObj)
        obj.write(protocol)
        return transObj.getvalue()

    # Implement the CReadableTransport interface.
    @property
    def cstringio_buf(self):
        return self.__rbuf

    def cstringio_refill(self, prefix, reqlen):
        # self.__rbuf will already be empty here because fastbinary doesn't
        # ask for a refill until the previous buffer is empty.  Therefore,
        # we can start reading new frames immediately.
        while len(prefix) < reqlen:
          self.readFrame()
          prefix += self.__rbuf.getvalue()
        self.__rbuf = StringIO(prefix)
        return self.__rbuf
        